#!/usr/bin/env python3
import argparse
from pathlib import Path
import base64
import hashlib
import hmac
import os


def decode_base64(data):
    """补齐并解码base64"""
    padding = '=' * (-len(data) % 4)
    return base64.b64decode(data + padding)


def hash_hostname(hostname, salt):
    """计算known_hosts的哈希值"""
    return base64.b64encode(hmac.new(salt, hostname.encode(), hashlib.sha1).digest()).decode()


def parse_known_hosts_line(line):
    """解析known_hosts行，返回(类型, hash, salt)"""
    parts = line.strip().split()
    if len(parts) < 3:
        return None
    if parts[0] != '|1|':
        return None
    salt = decode_base64(parts[1])
    hash_val = parts[2]
    return (hash_val, salt)


def normalize_hostname(hostname):
    """如果带端口，保持 [host]:port 格式，否则原样"""
    if ':' in hostname and not hostname.startswith('['):
        # 带端口的host要加中括号
        if hostname.count(':') == 1:
            host, port = hostname.split(':')
            return f"[{host}]:{port}"
        else:
            # ipv6地址可能有多个冒号，且可能本身带方括号
            if hostname.startswith('['):
                return hostname
            else:
                return hostname  # 保守处理
    return hostname


def main():
    parser = argparse.ArgumentParser(
        description='Match hashed known_hosts entries')
    parser.add_argument('hostname', help='hostname or [hostname]:port')
    parser.add_argument('--file', type=Path, default=Path.home() /
                        '.ssh' / 'known_hosts', help='known_hosts file path')
    args = parser.parse_args()

    hostname = normalize_hostname(args.hostname)
    kh_path = args.file.expanduser()

    if not kh_path.is_file():
        print(f"Known_hosts file not found: {kh_path}")
        return

    with kh_path.open() as f:
        for line in f:
            parsed = parse_known_hosts_line(line)
            if not parsed:
                continue
            hash_val, salt = parsed
            if hash_hostname(hostname, salt) == hash_val:
                print(f"Found matching hash line:\n{line.strip()}")
                return
    print("No matching hash found.")


if __name__ == '__main__':
    main()
